import cv2
import numpy as np
import pandas as pd
import math
import json
import copy
import maxflow as mf
from matplotlib import pyplot as plt
from scipy.signal import convolve2d
from sklearn.ensemble import AdaBoostClassifier
from sklearn import metrics
from operator import itemgetter

def load_annotations(file):
    with open(file) as json_file:
        annotations = json.loads(json_file.read())
    return annotations

def draw_annotations(img, annotation, color, outline_width):
    key = next(iter(annotation))
    regions = annotation[key]["regions"]
    for region in regions:
        cx = regions[region]["shape_attributes"]["cx"]
        cy = regions[region]["shape_attributes"]["cy"]
        r = regions[region]["shape_attributes"]["r"]
        cv2.circle(img, (cx,cy), r, color, outline_width)
    return img

def resize_annotations(annotations, downscale):
    key = next(iter(annotations))
    regions = annotations[key]["regions"]
    for region in regions:
        cx = int(regions[region]["shape_attributes"]["cx"] / downscale)
        cy = int(regions[region]["shape_attributes"]["cy"] / downscale)
        r = int(regions[region]["shape_attributes"]["r"] / downscale)
        annotations[key]["regions"][region]["shape_attributes"]["cx"] = cx
        annotations[key]["regions"][region]["shape_attributes"]["cy"] = cy
        annotations[key]["regions"][region]["shape_attributes"]["r"] = r
    return annotations

annotations = resize_annotations(annotations, 3.33333333333)

def split_mark(annotations, train_percent):
    regions = annotations["austrailian-mango-plantation.jpg2426773"]["regions"]
    sorted_regions = []
    for region in regions:
        lowerbound = regions[region]["shape_attributes"]["cy"] + regions[region]["shape_attributes"]["r"]
        sorted_regions.append((region, lowerbound))
    sorted_regions.sort(key=lambda tup: tup[1])
    return sorted_regions[int(len(sorted_regions)*train_percent)-1][1]

def train_test_split(img, mark, mode):
    if mode == 0:
        return (img[:mark, :], img[mark:, :])
    elif mode == 1:
        return (img[:, :mark], img[:, mark:])

lab = cv2.cvtColor(train, cv2.COLOR_BGR2LAB)
lab_features = np.reshape(lab, (-1, 3))

def bias_black(rgb_vec):
    rgb_vec[np.sum(rgb_vec, axis=1)==0] = [1,0,0]
    return rgb_vec

def bgr2xyz_vec(bgr):
    m = [[0.412453, 0.357580, 0.180423],
         [0.212671, 0.715160, 0.072169],
         [0.019334, 0.119193, 0.950227]]
    m = np.array(m)
    
    rgb = cv2.cvtColor(bgr, cv2.COLOR_BGR2RGB)
    rgb = np.reshape(rgb, (-1,3))
    rgb = bias_black(rgb)
    rgb = np.transpose(rgb)
    return np.true_divide(np.matmul(m, rgb), 2.55)

def illumination_invariant_colorspace_features(bgr):
    A = [[27.07439, -22.80783,-1.806681],
         [-5.646736, -7.722125, 12.86503],
         [-4.163133, -4.579428, -4.576049]]
    B = [[0.9465229, 0.2946927, -0.1313419],
         [-0.1179179, 0.9929960, 0.007371554],
         [0.09230461, -0.04645794, 0.9946464]]
    A = np.array(A)
    B = np.array(B)

    xyz_vec = bgr2xyz_vec(bgr)
    ii_features = np.transpose(np.matmul(A, np.log(np.matmul(B, xyz_vec))))
    return ii_features

def oriented_gaussian_derivative_filter2(image, sigma=1, theta=0, truncate=4.0):
    radius = int(truncate * sigma + 0.5)
    x = np.arange(-radius, radius+1)
    xx, yy = np.meshgrid(x, x)
    g0 = np.exp(-(xx**2 + yy**2)/(2*sigma**2)) / (sigma*math.sqrt(2*math.pi))
    G2a = -g0 / sigma**2 + np.multiply(g0, np.power(xx,2)) / sigma**4
    G2b = np.multiply(np.multiply(g0, xx),yy) / sigma**4
    G2c = -g0 / sigma**2 + np.multiply(g0,np.power(yy,2)) / sigma**4

    I2a = convolve2d(image, G2a, mode='same', boundary='symm')
    I2b = convolve2d(image, G2b, mode='same', boundary='symm')
    I2c = convolve2d(image, G2c, mode='same', boundary='symm')

    return np.cos(theta)**2*I2a + np.sin(theta)**2*I2c - 2*np.cos(theta)*np.sin(theta)*I2b

def get_texture_features(L):
    texture_features = None
    sigma = [1, math.sqrt(2), 2]
    orientations = np.linspace(0, math.pi, num=7, endpoint=True)[0:-1]
    for s in sigma:
        for o in orientations:
            texture = oriented_gaussian_derivative_filter2(L, sigma=s, theta=o)
            texture = np.reshape(texture, (-1, 1))
            if texture_features is not None:
                texture_features = np.concatenate((texture_features, texture), axis=1)
            else:
                texture_features = texture
    return texture_features

def entropy(signal):
    '''
    function returns entropy of a signal.
    signal must be a 1-D numpy array
    '''
    lensig = signal.size
    symset = list(set(signal))
    numsym = len(symset)
    propab = [np.size(signal[signal==i])/(1.0*lensig) for i in symset]
    ent = np.sum([p*np.log2(1.0/p) for p in propab])
    return ent

def get_entropy_features(L):
    N_list = [2, 4, 8]
    row, col = L.shape
    E = np.array(L)
    E_features = None
    for N in N_list:
        for r in range(row):
            for c in range(col):
                Lx = np.max([0,c-N])
                Ux = np.min([col,c+N+1])
                Ly = np.max([0,r-N])
                Uy = np.min([row,r+N+1])
                region = L[Ly:Uy, Lx:Ux].flatten()
                E[r, c] = entropy(region)
        E = np.reshape(E, (-1, 1))
        if E_features is not None:
            E_features = np.concatenate((E_features, E), axis=1)
        else:
            E_features = E
        E = np.array(L)
    return E_features

def is_pixel_tree(row):
    if row[0] == 0 and row[1] == 0 and row[2] == 0:
        return 1
    else:
        return 0

annotated_img = draw_annotations(img.copy(), annotations, (0,0,0), -1) 
annotated_img = annotated_img[:, :train_mark]
matrix = pd.DataFrame(np.reshape(annotated_img, (-1, 3)))
train_data["Tree"] = matrix.apply(is_pixel_tree, axis=1)

X_train = train_data.loc[:, train_data.columns != "Tree"]
y_train = train_data["Tree"]
abc = AdaBoostClassifier(n_estimators=200, learning_rate=1)
model = abc.fit(X_train, y_train)

y_pred = model.predict(X_train)

def V_p_q(label1, label2):
    return abs(label1-label2)

def D_p(label, proba, x, y):
    if label == 1:
        return math.log(1-proba[y][x])
    else:
        return math.log(proba[y][x])

def give_neighbours(image, x, y):
    if x>=len(image[0]) or x<0 or y>=len(image) or y<0:
       raise ValueError('Pixel is not in image. x and/or y are to large')
    ns = []
    for a,b in zip([1,0,-1,0],[0,1,0,-1]):
        if (x+a<len(image[0]) and x+a>=0) and (y+b<len(image) and y+b>=0):
            ns.append(image[y+b][x+a])
    return ns 

def return_mapping_of_image(image, alpha, beta):
    map = {}
    revmap = {}
    map_parameter = 0
    for y in range(len(image)):
        for x in range(len(image[0])):
            map[map_parameter] = (y,x)
            revmap[(y,x)] = map_parameter
            map_parameter += 1
    
    return map, revmap

def alpha_beta_swap_new(alpha, beta, tree, tree_proba):
    #extract position of alpha or beta pixels to mapping 
    map, revmap = return_mapping_of_image(tree, alpha, beta)
    
    #graph of maxflow 
    graph_mf = mf.Graph[float](len(map) )
    #add nodes
    nodes = graph_mf.add_nodes(len(map))
            
    #add n-link edges
    weight = 1
    for i in range(0,len(map)):
        y,x = map[i]
        #top, left, bottom, right
        for a,b in zip([1,0,-1,0],[0,1,0,-1]):
            if (y+b, x+a) in revmap:
                graph_mf.add_edge(i,revmap[(y+b,x+a)], weight, 0)
   
    #add all the terminal edges
    for i in range(0,len(map)):
        y,x = map[i]
        #find neighbours
        neighbours = give_neighbours(tree, x, y)
        #calculation of weight
        t_weight_alpha = sum([V_p_q(alpha,v) for v in neighbours]) + D_p(alpha, tree_proba, x, y)
        t_weight_beta = sum([V_p_q(beta,v) for v in neighbours]) + D_p(beta, tree_proba, x, y)
        graph_mf.add_tedge(nodes[i], t_weight_alpha, t_weight_beta)

    #calculating flow
    flow = graph_mf.maxflow()
    res = [graph_mf.get_segment(nodes[i]) for i in range(0, len(nodes))]
    
    #depending on cut assign new label
    for i in range(0, len(res)):
        y, x = map[i] 
        if res[i] == 1:
            tree[y][x] = alpha 
        else:
            tree[y][x] = beta
    
    return tree

def swap_minimization(tree, tree_proba, cycles):
    for i in range(0, cycles):
        tree  = alpha_beta_swap_new(0, 1, tree, tree_proba)     
        #user output and interims result image
        #print(str(u+1) + "\t\t\t", calculate_energy(img_orig, img_work)) 
        # print("Energy after " + str(u+1) + "/" + str(cycles) + " cylces:", calculate_energy(img_orig, img_work)) 
    
    return tree

y_pred_proba = model.predict_proba(X_train)[:,1]

tree_proba = np.reshape(y_pred_proba, (row, col))
tree = np.reshape(y_pred, (row, col))

refined_class = swap_minimization(tree, tree_proba, 10)

regions = annotations["austrailian-mango-plantation.jpg2426773"]["regions"]
training_regions = []
for region in regions:
    radius = regions[region]["shape_attributes"]["r"]
    cx = regions[region]["shape_attributes"]["cx"]
    cy = regions[region]["shape_attributes"]["cy"]
    upperbound = cy - radius
    lowerbound = cy + radius
    leftbound = cx - radius
    rightbound = cx + radius
    training_regions.append((upperbound, lowerbound, leftbound, rightbound, cx, cy, radius))
training_regions.sort(key=lambda tup: tup[2])

img_row, img_col, img_spec = img.shape

training_regions = [r for r in training_regions if (r[0] >= 0 and r[1] < img_row and r[2] >= 0 and r[3] < train_mark)]

training_regions = np.array([list(region) for region in training_regions])
radius = training_regions[:, 6]

def round_off_radius(regions, ro_radiuses, minmax):
    ro_regions = copy.deepcopy(regions)
    for i in range(len(regions)):
        region = regions[i]
        r = region[2]
        for ro in ro_radiuses:
            if r >= ro - minmax[0] and r <= ro + minmax[1]:
                r = ro
                break
        ro_regions[i][2] = r
    return ro_regions

ro_regions = round_off_radius(training_regions[:,4:], (7, 10, 13, 16, 19, 22, 25), (1, 1))
ro_regions = ro_regions[ro_regions[:,2].argsort()]
ro_regions = ro_regions[:-2]

def get_portion(img, cx, cy, r):
    row, col, spec = img.shape
    Lx = cx-r
    Ux = cx+r+1
    Ly = cy-r
    Uy = cy+r+1
    portion = img[Ly:Uy, Lx:Ux]
    return portion

def scale_region(region):
    scaled = region.copy()
    scaled = scaled.astype('float64')
    for i in range(len(scaled)):
        for j in range(len(scaled[i])):
            scaled[i][j] = np.true_divide(scaled[i][j], 255)
    return scaled

def is_pixel_P(row):
    if row[0] == 0 and row[1] == 0 and row[2] == 0:
        return 1
    else:
        return 0

def add_p(radius, scaled_average):
    row, col, spec = scaled_average.shape
    sa_copy = scaled_average.copy()
    cv2.circle(sa_copy, (int(row/2),int(col/2)), radius, 0, -1)
    scaled_average = pd.DataFrame(np.reshape(scaled_average, (-1, 3)))
    sa_copy = pd.DataFrame(np.reshape(sa_copy, (-1, 3)))
    scaled_average[4] = sa_copy.apply(is_pixel_P, axis=1)
    return np.reshape(scaled_average.values, (row, col, spec+1))

def create_templates(regions):
    templates = {}
    scaled_templates = []
    for region in regions:
        cx = region[0]
        cy = region[1]
        r = region[2]
        tree_portion = get_portion(img.copy(), cx, cy, r)
        if r not in templates:
            templates[r] = [tree_portion]
        else:
            templates[r].append(tree_portion)
    
    for template in templates:
        curr_template = templates[template]
        scaled_average = None
        num_of_regions = 0
        for region in curr_template:
            if scaled_average is None or scaled_average.shape < region.shape:
                scaled_average = scale_region(region)
                num_of_regions = 1
            elif scaled_average.shape == region.shape:
                scaled_regions = scale_region(region)
                scaled_average = np.add(scaled_average, scaled_regions)
                num_of_regions += 1
        scaled_average = np.true_divide(scaled_average, num_of_regions)
        scaled_average = add_p(template, scaled_average)
        scaled_templates.append((template, scaled_average)) 
    return scaled_templates

templates = create_templates(ro_regions)

row, col, spec = train.shape
target = scale_region(train)
target = np.reshape(target, (-1,3)).astype('float64')
target = np.concatenate((target, np.reshape(y_pred_proba, (-1, 1))), axis=1)
refined_pred = np.reshape(refined_class, (-1,3))
target = np.concatenate((target, np.reshape(refined_pred, (-1, 1))), axis=1)
target = np.reshape(target, (row, col, spec+2))

def tm_ccorr_normed(t, s):
    s, sx, sy, r = s
    s = s[sy-r:sy+1+r, sx-r:sx+1+r]
    row,col,spec = t.shape

    score = np.array([0, 0, 0, 0])
    for i in range(row):
        for j in range(col):
            if t[i][j][3] == 1:
                score = np.add(score, np.multiply(t[i][j][0:4], s[i][j][0:4]))

    t_energy = np.array([0, 0, 0, 0])
    s_energy = np.array([0, 0, 0, 0])
    for i in range(row):
        for j in range(col):
            if t[i][j][3] == 1:
                t_energy = np.add(t_energy, np.power(t[i][j][0:4], 2))
                s_energy = np.add(s_energy, np.power(s[i][j][0:4], 2))
    energy = np.sqrt(np.multiply(t_energy, s_energy))
    
    i_score = np.true_divide(score, energy)
    score = np.average(i_score)
#     print("{}, ({},{}), {}, {}".format(r, sx, sy, i_score, score))
    return score

def within_constraint(img, constraint):
    img, x, y, r = img
    img = img[y-r:y+1+r, x-r:x+1+r]
    row, col, spec = img.shape
    c_mask = np.zeros((row,col))
    cv2.circle(c_mask, (r,r), r, 1, -1)
    
    within_area = 0
    total_area = 0
    for i in range(row):
        for j in range(col):
            if c_mask[i][j] == 1:
                total_area += 1
                if img[i][j][4] == 1:
                    within_area += 1
    return within_area / total_area >= constraint

def match_templates(templates, img, thresh):
    candidates = []
    for template in templates:
        r = template[0]
        t = template[1]
        srow, scol, sspec = img.shape
        sx = r
        sy = r
        candidate_set = []
        while True:
            score = None
            if within_constraint((img,sx,sy,r), 0.75):
                score = tm_ccorr_normed(t, (img,sx,sy,r))

            if score and score >= thresh:
                candidate_set.append((score, sx, sy, r))
                
            if sx == scol-1-r and sy == srow-1-r:
                break
                
            if sx == scol-1-r:
                sx = r
                sy = np.min([sy+(r//2), srow-1-r])
            else:
                sx = np.min([sx+(r//2), scol-1-r])
        candidates.append(candidate_set)
    return candidates

def templates_dict_to_list_mapping(templates):
    mapping = []
    for template in templates:
        mapping.append((template[0], template[1]))
    return mapping

templates = templates_dict_to_list_mapping(templates)
templates.sort(key=itemgetter(0), reverse=True)
matches = match_templates(templates, target, 0.25)

def distance(a, b):
    xa = a[0]
    ya = a[1]
    xb = b[0]
    yb = b[1]
    return math.sqrt(((xa-xb)**2) + ((ya-yb)**2))

def remove_overlap(tm, thresh):    
    tmc = []
    for t in tm:
        t.sort(key=itemgetter(0), reverse=True)
        tmc.extend(t)
    
    final = []
    while len(tmc) != 0:
        top = tmc.pop(0)
        final.append(top)
        no_overlap = []
        for t in tmc:
            rtop = top[3]
            rt = t[3]
            ctop = (top[1], top[2])
            ct = (t[1], t[2])
            overlap = (rtop + rt - distance(ctop, ct)) / np.min([rtop, rt])
            if overlap <= thresh:
                no_overlap.append(t)
        tmc = no_overlap
    return final

final_matches = remove_overlap(matches, 0.25)

def draw_matches(img, matches, color, outline_width):
    for match in matches:
        cx = int(match[1])
        cy = int(match[2])
        r = int(match[3])
        cv2.circle(img, (cx,cy), r, color, outline_width)
    return img

test_lab = cv2.cvtColor(test, cv2.COLOR_BGR2LAB)
test_lab_features = np.reshape(test_lab, (-1, 3))
test_ii_features = illumination_invariant_colorspace_features(test)

test_texture_features = get_texture_features(test_lab[:, :, 0])

test_entropy_features = get_entropy_features(test_lab[:, :, 0])

test_data = np.concatenate((test_lab_features, test_ii_features, test_texture_features, test_entropy_features), axis=1)

test_y_pred = model.predict(test_data)

test_annotated_img = draw_annotations(img.copy(), annotations, (0,0,0), -1)
test_annotated_img = test_annotated_img[:, test_mark:]

test_matrix = pd.DataFrame(np.reshape(test_annotated_img, (-1, 3)))
test_y_ground = test_matrix.apply(is_pixel_tree, axis=1)

test_y_pred_proba = model.predict_proba(test_data)[:,1]

test_tree_proba = np.reshape(test_y_pred_proba, (test_row, test_col))
test_tree = np.reshape(test_y_pred, (test_row, test_col))

test_refined_class = swap_minimization(test_tree, test_tree_proba, 10)